import argparse

import torch
import torch.backends.cudnn as cudnn
import numpy as np
import PIL.Image as pil_image
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader

from datasets import EvalDataset
from models import FSRCNN
from utils import convert_ycbcr_to_rgb, preprocess, calc_psnr, AverageMeter


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights-file', type=str, required=True)
    parser.add_argument('--image-file', type=str, required=True)
    parser.add_argument('--scale', type=int, default=3)
    args = parser.parse_args()

    cudnn.benchmark = True
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    model = FSRCNN(scale_factor=args.scale).to(device)

    state_dict = model.state_dict()
    for n, p in torch.load(args.weights_file, map_location=lambda storage, loc: storage).items():
        if n in state_dict.keys():
            state_dict[n].copy_(p)
        else:
            raise KeyError(n)

    model.eval()

    image = pil_image.open(args.image_file).convert('RGB')

    image_width = (image.width // args.scale) * args.scale
    image_height = (image.height // args.scale) * args.scale

    hr = image.resize((image_width, image_height), resample=pil_image.BICUBIC)
    lr = hr.resize((hr.width // args.scale, hr.height // args.scale), resample=pil_image.BICUBIC)
    bicubic = lr.resize((lr.width * args.scale, lr.height * args.scale), resample=pil_image.BICUBIC)
    # 裁剪hr符合新的模型
    hr = hr[(7, image_width-7),  (7, image_height-7)]

    bicubic.save(args.image_file.replace('.', '_bicubic_x{}.'.format(args.scale)))

    lr, _ = preprocess(lr, device)
    hr, _ = preprocess(hr, device)
    _, ycbcr = preprocess(bicubic, device)

    epoch_psnr_Set5 = AverageMeter()
    eval_dataset_Set5 = EvalDataset("BLAH_BLAH/Set5-11_19_with_bicubic.h5")
    eval_dataloader_Set5 = DataLoader(dataset=eval_dataset_Set5, batch_size=2)
    for data in eval_dataloader_Set5:
        inputs, labels, bicubic = data

        inputs = inputs.to(device)
        labels = labels.to(device)
        bicubic = bicubic.to(device)

        with torch.no_grad():
            preds = model(inputs)

        # preds = torch.squeeze(preds)
        # inputs = torch.squeeze(inputs)
        # labels = torch.squeeze(labels)
        # print(inputs.shape)
        # bicubic = F.interpolate(inputs, scale_factor=19/11, mode="bicubic", align_corners=True, recompute_scale_factor=True)

        preds = preds + bicubic

        psnr = calc_psnr(preds, labels)
        psnr_bicubic_label = calc_psnr(bicubic, labels)

        epoch_psnr_Set5.update(psnr, len(inputs))
    print('eval psnr on Set5: {:.2f}'.format(epoch_psnr_Set5.avg))

    epoch_psnr_Set14 = AverageMeter()
    eval_dataset_Set14 = EvalDataset("BLAH_BLAH/Set14-11_19_with_bicubic.h5")
    eval_dataloader_Set14 = DataLoader(dataset=eval_dataset_Set5, batch_size=2)
    for data in eval_dataloader_Set14:
        inputs, labels, bicubic = data

        inputs = inputs.to(device)
        labels = labels.to(device)
        bicubic = bicubic.to(device)

        with torch.no_grad():
            preds = model(inputs)

        # preds = torch.squeeze(preds)
        # inputs = torch.squeeze(inputs)
        # labels = torch.squeeze(labels)
        # print(inputs.shape)
        # bicubic = F.interpolate(inputs, scale_factor=19/11, mode="bicubic", align_corners=True, recompute_scale_factor=True)

        preds = preds + bicubic

        psnr = calc_psnr(preds, labels)
        psnr_bicubic_label = calc_psnr(bicubic, labels)

        epoch_psnr_Set5.update(psnr, len(inputs))
    print('eval psnr on Set14: {:.2f}'.format(epoch_psnr_Set5.avg))

    with torch.no_grad():
        preds = model(lr)
    preds = preds + bicubic  # residual
    psnr = calc_psnr(hr, preds)
    print('PSNR: {:.2f}'.format(psnr))

    preds = preds.mul(255.0).cpu().numpy().squeeze(0).squeeze(0)

    output = np.array([preds, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
    output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
    output = pil_image.fromarray(output)
    output.save(args.image_file.replace('.', '_fsrcnn_xvier_res.'))


    first_layer = model.first_part
    last_layer = model.last_part
    feature_map1 = first_layer.__getitem__(0).weight.cpu().clone()
    feature_map2 = last_layer.__getitem__(0).weight.cpu().clone()
    print("number of the first layer: ", len(feature_map1))
    print("number of the last layer: ", len(feature_map2))
    plt.figure(figsize=(10, 10))
    for i in range(0, len(feature_map1)):
        map1 = feature_map1[i]
        plt.subplot(8, 7, i + 1)
        plt.axis('off')
        plt.imshow(map1[0, :, :].detach(), cmap='gray')
    plt.savefig('./data/layer_features/fsrcnn-xavier/fist_layer.png')

    plt.figure(figsize=(10, 10))
    for i in range(0, len(feature_map2)):
        map2 = feature_map2[i]
        plt.subplot(8, 7, i + 1)
        plt.axis('off')
        plt.imshow(map2[0, :, :].detach(), cmap='gray')
    plt.savefig('./data/layer_features/fsrcnn-xavier/last_layer.png')
